/*
User-Assisted Intrinsic Images
Adrien Bousseau, Sylvain Paris, Fredo Durand.

This is the pseudo-code to generate the L matrix according to equation 11 of the paper,
and solve equation 16 with Gauss-Seidel iterations.
*/


//----- LMatrix initialization -----//

#define WSIZE 3 //size of the local windows
#define WSIZE2 9 //WSIZE2 = WSIZE*WSIZE, number of pixels in each window ('n' in the paper)
#define MSIZE 12 //WSIZE2+3, number of rows in the M matrix (see equation 8)
#define WSIZE_2 1 //WSIZE_2 = WSIZE/2
#define SIZE_N 5 //2*2*WSIZE_2+1, size of the total neighborhood affected by a pixel
#define SIZE_N2 25 //SIZE_N*SIZE_N, number of non-zero element per-pixel in the L matrix 
double ** LMatrix;

void buildLMatrix(Img imgIn)
{
	double epsilon = 0.0000001;
	int w = imgIn.width();
	int h = imgIn.height();

	double M[MSIZE][3]; //M matrix (see equation 8)
	M[MSIZE-3][0] = sqrt_epsilon; M[MSIZE-3][1] = 0; M[MSIZE-3][2] = 0;
	M[MSIZE-2][0] = 0; M[MSIZE-2][1] = sqrt_epsilon; M[MSIZE-2][2] = 0;
	M[MSIZE-1][0] = 0; M[MSIZE-1][1] = 0; M[MSIZE-1][2] = sqrt_epsilon;

	int i,j,ik,jk;
	int x,y;
	int xC,yC,xiC,yiC,xjC,yjC;

	//Init the L matrix. L is a sparse matrix that only contains SIZE_N2 non-zero elements per row.
	//We store the L matrix in an array with one row per image pixel and one column per non-zero value. 
	//Each array cell accumulates values according to equation 11.
	LMatrix = new double* [w*h];
	for(i=0; i<w*h; i++){
		LMatrix[i] = new double[SIZE_N2];
		for(j=0; j<SIZE_N2; j++){
			LMatrix[i][j] = 0;
		}
	}

	//Compute the L matrix
	for( y = WSIZE_2; y < h-WSIZE_2; y++) {//for each pixel...
		for( x = WSIZE_2; x < w-WSIZE_2; x++) {
			//fill the M matrix with the image values in the current window
			for (int m=-WSIZE_2; m<=WSIZE_2; m++) {
				yC = y+m;
				for (int n=-WSIZE_2; n<=WSIZE_2; n++) {
					xC = x+n;
					M[(m+WSIZE_2)*WSIZE + n+WSIZE_2][0] = imgIn.getR(xC,yC);
					M[(m+WSIZE_2)*WSIZE + n+WSIZE_2][1] = imgIn.getG(xC,yC);
					M[(m+WSIZE_2)*WSIZE + n+WSIZE_2][2] = imgIn.getB(xC,yC);
				}
			}
			//Compute the N matrix (equation 9 and 10)
			Matrix<MSIZE,3,double> Mi(M);
			Matrix<3,MSIZE,double> Mit(Mi.transpose());
			Square_matrix<3,double> MitMi(Mit*Mi);
			MitMi_inv = inverse_matrix(MitMi);
			Square_matrix<MSIZE,double> Ni(Mi*MitMi_inv*Mit);
			Ni = Ni.identity() - Ni;
			Square_matrix<MSIZE,double> Nit(Ni.transpose());
			Square_matrix<MSIZE,double> NitNi(Nit*Ni);

			//accumulate the values in the L matrix, for each couple of pixels (i,j) covered by the window
			//--- i
			for (int mi=-WSIZE_2; mi<=WSIZE_2; mi++) {//(mi,ni) = coordinate of pixel i in the local window
				yiC = y+mi; //(xiC,yiC) = coordinate of pixel i in the image
				for (int ni=-WSIZE_2; ni<=WSIZE_2; ni++) {
					xiC = x+ni;
					i = yiC*w+xiC; //row number in the matrix
					ik = (mi+WSIZE_2)*WSIZE + ni+WSIZE_2;
					//--- j
					for (int mj=-WSIZE_2; mj<=WSIZE_2; mj++) {
						yjC = y+mj;
						for (int nj=-WSIZE_2; nj<=WSIZE_2; nj++) {
							xjC = x+nj;
							jk = (mj+WSIZE_2)*WSIZE + nj+WSIZE_2;
							//accumulate values (equation 11)
							LMatrix[i][(yiC-yjC+2*WSIZE_2)*SIZE_N+(xiC-xjC+2*WSIZE_2)] += NitNi(ik,jk);
							//pixel j is at coordinate (xiC-xjC,yiC-yjC) in the SIZE_N*SIZE_N neighborhood affected by pixel i (centered at pixel i)
						}
					}
				}
			}
		}
	}
}



//----- Gauss-Seidel iterations -----//

/*
imgScribbleAbsolute: image that stores the fixed-illumination scribbles (pixels store illumination values, alpha = true if the pixel is scribbled)
imgScribbleSimAlbedo: image that stores the constant-reflectance scribbles (pixels store scribbles ids)
imgScribbleSimShading: image that stores the constant-illumination scribbles (pixels store scribbles ids)
maxAlbedoNum: number of constant-reflectance scribble ids
maxShadingNum: number of constant-illumination scribble ids
albedoNumColors: constant-reflectance scribble ids
shadingNumColors: constant-illumination scribble ids
*/
#define EPSILON_PIXEL_VAL 0.001 //accuracy threshold

Img solve(Img imgIn, Img &imgScribbleAbsolute,Img imgScribbleSimAlbedo,Img imgScribbleSimShading,int maxAlbedoNum, int maxShadingNum,
	Color* albedoNumColors, Color* shadingNumColors)
{
	int w = imgIn.width();
	int h = imgIn.height();
	Img result(w,h,Color(0.5,0.5,0.5));
	double rRes, gRes, bRes;//temporary result values
	double denom;
   	double lambdaCstr = 0.5;//weight of the user constraints

	//variables used to compute user constraints
	double rSumI, gSumI, bSumI;//sum of images values
    	double rSumS, gSumS, bSumS;//sum of illumination values
	double rAvrg, gAvrg, bAvrg;//average illumination
	Color imgVal, resVal;
        Color* albedoAvgColors = new Color[maxAlbedoNum];//average reflectance for each constant-reflectance scribble
        Color* shadingAvgColors = new Color[maxShadingNum];//average illumination for each constant-illumination scribble

	//Gauss-Seidel iterations.
        //for simplicity, we apply here a fixed number of iterations
	for (int k=0; k<nbIter; k++){
		//-------------------------------//
		//--- User constraints values ---//
		//-------------------------------//

		//In order to apply the user constraints, we first need to compute the average value of constant reflectance or illumination pixels
		//-- Constant-reflectance
		for (int c=0; c<=maxAlbedoNum; c++){//for each constant-reflectance scribble id
		    //illumination estimation for pixel p = I(p)*sum(S)/sum(I), where S stands for illumination values of constant pixels and I for input value of constant pixels
		    rSumI = gSumI = bSumI = 0;
		    rSumS = gSumS = bSumS = 0;
		    nbP = 0;
		    for (int y = 0; y < h; y++) {//for each pixel of the scribbles image
		        for (int x = 0; x < w; x++) {
		            if(scribbleSimAlbedo.pixel(x,y)==albedoNumColors[c]){//pixel scribbled with the current constraint
		                //get the image values for this pixel
		                imgVal = imgIn.pixel(x,y);
		                resVal = result.pixel(x,y);
		                rSumI += imgVal.getR();
		                rSumS += resVal.getR();
		                gSumI += imgVal.getG();
		                gSumS += resVal.getG();
		                bSumI += imgVal.getB();
		                bSumS += resVal.getB();
		                nbP++;
		            }
		        }
		    }
		    if(nbP>0){//if there was some scribbled pixels
		        albedoAvgColors[c] = Color(rSumS/rSumI,gSumS/gSumI,bSumS/bSumI);
		    }
		}
		//-- Constant-illumination
		for (int c=0; c<=maxShadingNum; c++){//for each constant-illumination scribble id
		    //illumination estimation = average illumination
		    rAvrg = gAvrg = bAvrg = 0;
		    nbP = 0;
		    for (int y = 0; y < h; y++) {//for each pixel of the scribbles image
		        for (int x = 0; x < w; x++) {
		            if(scribbleSimShading.pixel(x,y)==shadingNumColors[c]){//pixel scribbled with the current constraint
		                resVal = result.pixel(x,y);
		                rAvrg += resVal.getR();
		                gAvrg += resVal.getG();
		                bAvrg += resVal.getB();
		                nbP++;
		            }
		        }
		    }
		    if(nbP>0){//if there was some scribbled pixels
		        rAvrg/=nbP;
		        gAvrg/=nbP;
		        bAvrg/=nbP;
		        shadingAvgColors[c] = Color(rAvrg,gAvrg,bAvrg);
		    }
		}

		//------------------------------------------------//
		//--- Gauss-Seidel iteration with the L matrix ---//
		//------------------------------------------------//

		for (int y = 0; y < h; y++) {
		    for (int x = 0; x < w; x++) {
		        i = y*w+x; //row number for the current pixel
		        //is that a fixed-illumination pixel?
		        imgVal = scribbleAbsolute.pixel(x,y);
		        if(imgVal.getA()){//fixed-illumination pixel: impose the constraint
		            rRes = imgVal.getR();
		            gRes = imgVal.getG();
		            bRes = imgVal.getB();
		        }
		        else{//non scribbled pixel
		            //-- Gauss-Seidel update
		            rRes = gRes = bRes = 0;
		            for (int k=0; k < SIZE_N2; k++){//for each value on the matrix row corresponding to the current pixel
		                if(LMatrix[i][k]!=0){
					//coordinate of the corresponding pixel in the image
		                    yjC = -k/SIZE_N+2*WSIZE_2+y;
		                    xjC = -k+(y-yjC+2*WSIZE_2)*SIZE_N+2*WSIZE_2+x;

		                    if(x==xjC && y==yjC){//diagonal
		                        denom = LMatrix[i][k];
		                    }
		                    else{
		                        resVal = result.pixel(xjC,yjC);
		                        //relaxation
		                        rRes -= LMatrix[i][k]*resVal.getR();
		                        gRes -= LMatrix[i][k]*resVal.getG();
		                        bRes -= LMatrix[i][k]*resVal.getB();
		                    }
		                }
		            }
		            rRes /= denom;
		            gRes /= denom;
		            bRes /= denom;
		            //-- User constraints updates
		            //check for constant-reflectance
		            scribVal = scribbleSimAlbedo.pixel(x,y);
		            if(scribVal.getA()){//constant-reflectance pixel
		                //get the scribble id
		                int id = albedoNumColors.indexOf(QColor(scribVal));
		                //get the image value for this pixel
		                imgVal = imgIn.pixel(x,y);
		                //illumination = I(p) * sum(S)/sum(I)
		                //if one channel is too small, take another one
		                //red
		                if(imgVal.getR()>EPSILON_PIXEL_VAL && albedoAvgColors[id].getR()>EPSILON_PIXEL_VAL)
		                    rRes = lambdaCstr*rRes + (1.0-lambdaCstr)*imgVal.getR()*albedoAvgColors[id].getR();
		                else{
		                    if(imgVal.getG()>EPSILON_PIXEL_VAL && albedoAvgColors[id].getG()>EPSILON_PIXEL_VAL)
		                        rRes = lambdaCstr*gRes + (1.0-lambdaCstr)*imgVal.getG()*albedoAvgColors[id].getG();
		                    else
		                        rRes = lambdaCstr*bRes + (1.0-lambdaCstr)*imgVal.getB()*albedoAvgColors[id].getB();
		                }
		                //green
		                if(imgVal.getG()>EPSILON_PIXEL_VAL && albedoAvgColors[id].getG()>EPSILON_PIXEL_VAL)
		                    gRes = lambdaCstr*gRes + (1.0-lambdaCstr)*imgVal.getG()*albedoAvgColors[id].getG();
		                else
		                    gRes = rRes;
		                //blue
		                if(imgVal.getB()>EPSILON_PIXEL_VAL && albedoAvgColors[id].getB()>EPSILON_PIXEL_VAL)
		                    bRes = lambdaCstr*bRes + (1.0-lambdaCstr)*imgVal.getB()*albedoAvgColors[id].getB();
		                else
		                    bRes = rRes;
		            }
		            else{
		                //check for constant-illumination
		                scribVal = scribbleSimShading[deepness-1].pixel(x,y);
		                if(qAlpha(scribVal)>1){//constant-illumination pixel
		                    //get the scribble id
		                    int id = shadingNumColors.indexOf(QColor(scribVal));
		                    rRes = lambdaCstr*rRes + (1.0-lambdaCstr)*shadingAvgColors[id].getR();
		                    gRes = lambdaCstr*gRes + (1.0-lambdaCstr)*shadingAvgColors[id].getG();
		                    bRes = lambdaCstr*bRes + (1.0-lambdaCstr)*shadingAvgColors[id].getB();
		                }
		            }
		        }
		        if(rRes<EPSILON_PIXEL_VAL)
		            rRes = EPSILON_PIXEL_VAL;
		        if(gRes<EPSILON_PIXEL_VAL)
		            gRes = EPSILON_PIXEL_VAL;
		        if(bRes<EPSILON_PIXEL_VAL)
		            bRes = EPSILON_PIXEL_VAL;
		        result.setPixel(x,y,Color(rRes,gRes,bRes));
		    }
		}
	}
}
